import matplotlib.pyplot as  plt
import numpy as np


def _numerical_gradient_no_batch(f, x):
    h = 1e-4  # 数值微分的思想求近似切线斜率
    grad = np.zeros_like(x)

    for idx in range(x.size):
        tmp_val = x[idx]
        x[idx] = float(tmp_val) + h
        fxh1 = f(x)  # f(x+h)
        x[idx] = tmp_val - h
        fxh2 = f(x)  # f(x-h)

        grad[idx] = (fxh1 - fxh2) / (2 * h)  # 以数值微分的思想计算梯度
    return grad


def numerical_gradient(f, X):
    if X.ndim == 1:
        return _numerical_gradient_no_batch(f, X)
    else:
        grad = np.zeros_like(X)  # 生成一个和参数.shape一致的数据全为0的矩阵

        for idx, x in enumerate(X):
            grad[idx] = _numerical_gradient_no_batch(f, x)

        return grad


def function_2(x):
    if x.ndim == 1:
        return np.sum(x ** 2)
    return np.sum(x ** 2, axis=1)


if __name__ == '__main__':
    x0 = np.arange(-2, 2.5, 0.25)
    x1 = np.arange(-2, 2.5, 0.25)

    X, Y = np.meshgrid(x0, x1)
    X = X.flatten()
    Y = Y.flatten()

    grad = numerical_gradient(function_2, np.array([X, Y]))

    # 使用plt画出剑鞘
    plt.quiver(X, Y, -grad[0], -grad[1], angles="xy", color="#666666")  # ,headwidth=10,scale=40,color="#444444")

    plt.title('function\'s gratitude Image')
    plt.xlabel('x0')
    plt.ylabel('x1')
    plt.xlim([-2, 2])
    plt.ylim([-2, 2])

    plt.grid()
    plt.show()
